package functionalEnrichment.semanticSimilarity.semanticSimilarityCalculator.program.geneSetAnalysis;

import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.CountDownLatch;
import org.apache.commons.collections.FastHashMap;
import org.apache.commons.math.stat.descriptive.SummaryStatistics;
import functionalEnrichment.semanticSimilarity.semanticSimilarityCalculator.parsers.SimilarityFileParser;
import functionalEnrichment.semanticSimilarity.semanticSimilarityCalculator.units.Gene;
import functionalEnrichment.semanticSimilarity.semanticSimilarityCalculator.units.GeneLibrary;
import functionalEnrichment.semanticSimilarity.semanticSimilarityCalculator.units.SemanticSimilarity;
import functionalEnrichment.semanticSimilarity.semanticSimilarityCalculator.utils.FileUtils;

/**
 * This class performs is invoked by the <code>GeneSetAnalysisProgram</code> class.
 * It does the main workin the GeneSetAnalysis program.
 * A summary of the semantic similarities between gene sets that are passed in to the constructor are
 * calculated in an all-against-all comparison. A summary of the similarities between gene sets
 * are outputted in to the results file that is specified on the command line.
 * This program will run in multiple threads, where gene set comparisons are split in to roughly equal
 * sets. This  feature that is required if the
 * number and size of the gene sets is very large, as all-against-all comparisons
 * can involve a very large number of gene-to-gene similarities.
 * @author Jamie MacPherson
 */
class GeneSetAnalysis {

    /*
     *The gene sets are contained in the lists of genes and stored in a map
    with a string key (the identifier for that gene set.
     */
    private final Map<String, List<Gene>> geneSets;
    // The gene sets are filtered to include only annotated genes, also stored in a map.
    private Map<String, List<Gene>> filteredGeneSets;
    // the semantic similarities between genes are stored in a map
    private final FastHashMap similarityMap;
    // the gene library,  a repository for all gene objects
    private final GeneLibrary geneLibrary;
    // a list of the annotated genes
    private final List<Gene> allAnnotatedGenes;
    // the number of threads that the program should run in.
    private final int noOfThreads;
    // a folder that the results will be written to.
    private final File resultsFile;
    // if set to true, gene sets will only be compared against themselves
    private boolean runInternalOnly = false;

    /**
     * Instantiates the class
     * @param geneSets The gene sets are contained in the lists of genes and stored in a map
     * with a string key (the identifier for that gene set.)
     * @param similarityFileParser contains all of the data for gene-to-gene semantic similarities, annotated genes and a gene libraray.
     * @param noOfThreads the number of threads that the program will run in
     * @param resultsFile a folder that the resulst will be written to.
     */
    GeneSetAnalysis(Map<String, List<Gene>> geneSets, SimilarityFileParser similarityFileParser, int noOfThreads, File resultsFile) {
        this.geneSets = geneSets;
        this.geneLibrary = similarityFileParser.getGeneLibrary();
        this.allAnnotatedGenes = similarityFileParser.getAnnotatedGenes();
        this.noOfThreads = noOfThreads;
        this.resultsFile = resultsFile;
        similarityMap = similarityFileParser.getSimilarityHash();
    }

    /**
     * After instantiation of the class, this method will run the gene set analysis.
     */
    void run() {
        Collections.sort(allAnnotatedGenes);
        filteredGeneSets = getFilteredGeneSets();
        Runner[] runners = new Runner[noOfThreads];
        System.out.println("the gene set analysis program will run in " + noOfThreads + " threads");
        CountDownLatch cdl = new CountDownLatch(noOfThreads);
        createRunners(runners, cdl);
        for (Runner runner : runners) {
            runner.start();
        }
        try {
            cdl.await();
        } catch (InterruptedException ex) {
            ex.printStackTrace();
        }
    }

    /**
     * Filters the gene sets map and returns an equivalent map (with the same keyset)
     * that only includes genes annotated with informative GO terms.
     * @return the map of filtered gene lists.
     */
    private Map<String, List<Gene>> getFilteredGeneSets() {
        Map<String, List<Gene>> filtered = new HashMap<String, List<Gene>>();
        for (String geneSetId : geneSets.keySet()) {
            filtered.put(geneSetId, filterList((List<Gene>) geneSets.get(geneSetId)));
        }
        return filtered;

    }

    /**
     * Filters a list of genes to only include genes annotated with informative GO terms.
     * @param geneList a list of genes to filter.
     * @return a filtered gene list.
     */
    private List<Gene> filterList(List<Gene> geneList) {
        List<Gene> filtered = new ArrayList<Gene>();
        for (Gene gene : geneList) {
            int index = Collections.binarySearch(allAnnotatedGenes, gene, null);
            if (index >= 0) {
                filtered.add(gene);
            }
        }
        return filtered;
    }

    /**
     * Given an array of <code>Runner</code> objects (that is equal in length
     * to the number of threads the program will run in, this method will instantiate
     * the Runners and assign a roughly equal number of gene sets to each. Each runner must have reference
     * to a the same <code>CountDownLatch</code> object, in order to notify the program when the runner has completed it's analysis.
     * @param runners an array of runners with a length equal to the number of threads the program will run in.
     * @param countDownLatch to count down when the runner is finished.
     */
    private void createRunners(Runner[] runners, CountDownLatch countDownLatch) {
        Map<Integer, List<GeneSetPair>> split = getThreadSplit();

        for (int i = 0; i < split.keySet().size(); i++) {
            runners[i] = new Runner(i, split.get(i), countDownLatch);
        }
    }

    /**
     * A method that will split all of the gene set comparisons between a number of groups equal in size
     * to the number of threads in which the program will run. Each list of gene set pairs
     * will be run in a different thread by the program.
     * Returns a map of <code>GeneSetPair</code> list objects split up in a hash.
     * @return
     */
    private Map<Integer, List<GeneSetPair>> getThreadSplit() {
        Map<Integer, List<GeneSetPair>> split = new TreeMap<Integer, List<GeneSetPair>>();
        int splitId = 0;
        String[] keySet = filteredGeneSets.keySet().toArray(new String[filteredGeneSets.keySet().size()]);
        if (!runInternalOnly) {
            for (int i = 0; i < keySet.length; i++) {
                String gsId1 = keySet[i];
                for (int j = i; j < keySet.length; j++) {
                    String gsId2 = keySet[j];
                    if (!split.containsKey(splitId)) {
                        split.put(splitId, new ArrayList<GeneSetPair>());
                    }
                    GeneSetPair gsp = new GeneSetPair(gsId1, gsId2);
                    split.get(splitId).add(gsp);
                    splitId++;
                    if (splitId == noOfThreads) {
                        splitId = 0;
                    }
                }
            }
        } else {
            for (int i = 0; i < keySet.length; i++) {
                String gsId1 = keySet[i];
                String gsId2 = keySet[i];
                if (!split.containsKey(splitId)) {
                    split.put(splitId, new ArrayList<GeneSetPair>());
                }
                GeneSetPair gsp = new GeneSetPair(gsId1, gsId2);
                split.get(splitId).add(gsp);
                splitId++;
                if (splitId == noOfThreads) {
                    splitId = 0;
                }
            }
        }
        return split;
    }

    /**
     * A runner is a Thread object that receives the data to analyse when instantiated.
     * The <code>run()</code> method of the thread object has been overridden, this method
     * will run the analysis. After the analysis has run, the <code>.countDown()</code> method
     * of CountDownLatch object passed in to the constructor is called in order to notify the
     * SemanticSimilarityCalculator that the analysis is finished.
     * Results of the analysis are written to a new file in the results folder.
     */
    private class Runner extends Thread {

        // a list of gene set pairs - gene sets that will be analysed for semantic similarity.
        private List<GeneSetPair> geneSetPairs;
        // will count down when the analysis is complete to notify the program.
        private CountDownLatch cdl;
        // the writer to which the results are written.
        private BufferedWriter writer;

        /**
         * Instantiate a runner object. This requires an ID in order to write to a unique output file,
         * a set of <code>GeneSetPairs</code> that will be analysed for semantic similarity and a
         * <code>CounDownLatch</code> that will count down when the analysis is complete to notify the program.
         * @param runnerId the ID for this runner
         * @param geneSetPairs that will be analysed for semantic similarity
         * @param countDownLatch that will count down when the analysis is complete to notify the program.
         */
        public Runner(int runnerId, List<GeneSetPair> geneSetPairs, CountDownLatch countDownLatch) {
            this.geneSetPairs = geneSetPairs;
            this.cdl = countDownLatch;
            writer = FileUtils.getBufferedWriter(new File(resultsFile, "output_" + runnerId + ".txt"));
        }

        /**
         * Overrides the <code>run()</code> method of the <code>Thread</code> class.
         * In this method, the analysis is run -- an all-against-all analysis of gene the sets.
         * Results of the average semantic similarity, standard deviation of the semantic similarity,
         * the number of genes in each set and the number of genes in each set that are annotated by
         * informative GO terms.
         */
        @Override
        public void run() {
            System.out.println(geneSetPairs.size() + " gene sets in this runner");
            for (int i = 0; i < geneSetPairs.size(); i++) {
                String id1 = geneSetPairs.get(i).getId1();
                String id2 = geneSetPairs.get(i).getId2();
                List<Gene> gs1 = (List<Gene>) filteredGeneSets.get(id1);
                List<Gene> gs2 = (List<Gene>) filteredGeneSets.get(id2);
                SummaryStatistics ss = getSimilarityStats(gs1, gs2);
                Result r = new Result(id1, id2);
                r.setGs1_size(((List<Gene>) geneSets.get(id1)).size());
                r.setGs2_size(((List<Gene>) geneSets.get(id2)).size());
                r.setGs1_filteredSize(gs1.size());
                r.setGs2_filteredSize(gs2.size());
                r.setMeanSimilarity(ss.getMean());
                r.setSd(ss.getStandardDeviation());
                try {
                    writer.append(r.toString() + "\n");
                    writer.flush();
                } catch (IOException ex) {
                    ex.printStackTrace();
                }
            }
            try {
                writer.close();
            } catch (IOException ex) {
                ex.printStackTrace();
            }
            cdl.countDown();

        }

        /**
         * Given two gene lists, this method will return a <code>SummaryStatistics</code> object that
         * contains a summary of the semantic similarity between the two gene sets.
         * @param geneList1 the first gene list
         * @param geneList2 the second gene list
         * @return a <code>SummarStatistics</code> object that has the statistics for the semantic similarity
         * between the two gene sets.
         */
        private SummaryStatistics getSimilarityStats(List<Gene> geneList1, List<Gene> geneList2) {
            SummaryStatistics stats = new SummaryStatistics();
            for (Gene g1 : geneList1) {
                for (Gene g2 : geneList2) {
                    String[] ids = {g1.getId(), g2.getId()};
                    boolean add = true;
                    if (ids[0].equals(ids[1]) && runInternalOnly) {
                        add = false;
                    }
                    if (add) {
                        Arrays.sort(ids);
                        SemanticSimilarity ss = new SemanticSimilarity(geneLibrary.getGene(ids[0]), geneLibrary.getGene(ids[1]));
                        int v = 0;
                        if (similarityMap.containsKey(ss.hashCode())) {
                            v = ((SemanticSimilarity) similarityMap.get(ss.hashCode())).getSimilarity();
                        }
                        stats.addValue((double) v);
                    }

                }
            }
            return stats;
        }
    }

    public boolean isRunInternalOnly() {
        return runInternalOnly;
    }

    public void setRunInternalOnly(boolean runInternalOnly) {
        this.runInternalOnly = runInternalOnly;
    }
}
